import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

class Plotter:
    def __init__(self, subchannel):
        self.subchannel = subchannel
        
    def plot_axial_distribution(self, variable, title, ylabel):
        """绘制轴向分布图"""
        plt.figure(figsize=(10, 6))
        z = np.linspace(0, 1, self.subchannel.n_axial)
        
        for j in range(self.subchannel.n_channels):
            plt.plot(z, variable[:, j], label=f'Channel {j+1}')
            
        plt.xlabel('Normalized Axial Position')
        plt.ylabel(ylabel)
        plt.title(title)
        plt.legend()
        plt.grid(True)
        plt.show()
        
    def plot_void_fraction_3d(self):
        """绘制空泡率3D分布图"""
        fig = plt.figure(figsize=(12, 8))
        ax = fig.add_subplot(111, projection='3d')
        
        z = np.linspace(0, 1, self.subchannel.n_axial)
        channels = np.arange(self.subchannel.n_channels)
        Z, C = np.meshgrid(z, channels)
        
        surf = ax.plot_surface(Z, C, self.subchannel.void_fraction.T,
                             cmap='viridis', edgecolor='none')
        
        ax.set_xlabel('Normalized Axial Position')
        ax.set_ylabel('Channel Number')
        ax.set_zlabel('Void Fraction')
        ax.set_title('3D Void Fraction Distribution')
        
        fig.colorbar(surf)
        plt.show()
        
    def plot_convergence_history(self, mass_balance_history, energy_balance_history):
        """绘制收敛历史"""
        plt.figure(figsize=(12, 5))
        
        plt.subplot(1, 2, 1)
        plt.semilogy(mass_balance_history)
        plt.xlabel('Iteration')
        plt.ylabel('Mass Balance Error')
        plt.grid(True)
        
        plt.subplot(1, 2, 2)
        plt.semilogy(energy_balance_history)
        plt.xlabel('Iteration')
        plt.ylabel('Energy Balance Error')
        plt.grid(True)
        
        plt.tight_layout()
        plt.show() 